import pydevd_pycharm
import datetime
import functools

from absl import app, flags

import coin.proto.archive_flag_pb2 as afpb
import coin.proto.coin_feed_pb2 as feedpb
import coin.strategy.mm.tool.archive_base as abase
from coin.base.datetime_util import to_timestamp_int
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
from coin.strategy.mm.tool.archive_base import run_from_feed_cache_base


class FeedReplayer(object):
  def __init__(self, products, start_time=None, end_time=None):
    self._products = products
    self._start_ts = to_timestamp_int(start_time) or 0
    self._end_ts = to_timestamp_int(end_time) or 0
    self._feed = []

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, 'subscribe_trade'):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_book(self, product, book):
    bid0p = book.bid0().price
    bid0q = book.bid0().qty
    ask0p = book.ask0().price
    ask0q = book.ask0().qty
    midp = (bid0p + ask0p) / 2 if bid0p is not None and ask0p is not None else None
    timestamp = book.timestamp
    self._feed.append({
        'timestamp': timestamp,
        'bid0p': bid0p,
        'bid0q': bid0q,
        'ask0p': ask0p,
        'ask0q': ask0q,
        'midp': midp
      })

  def on_trade(self, product, trade):
    self._feed.append({
        'timestamp': trade.timestamp,
        'trade_side': feedpb.TradeSide.Name(trade.side),
        'tradep': trade.price,
        'tradeq': trade.qty,
      })


def main(argv):
  feed_machine = 'feed-best.ap-northeast-1.aws'
  feed_cache_dir = '/remote/iosg/coin-2/feed_cache'
  worker_ids = '1,2'
  start_time = datetime.datetime(2020, 8, 5)
  end_time = datetime.datetime(2020, 8, 5, 1, 0, 0)
  products = [generate_product_from_str2(
      'Futures', 'Huobi', None, 'BTC-USD.20200925', start_time)]

  fs_proto = afpb.FeedSubscriptionFlagProto(
      feed_machine=feed_machine,
      worker_ids=worker_ids,
      use_feed_cache=True,
      feed_cache_dir=feed_cache_dir)
  replayer = FeedReplayer(products, start_time, end_time)

  pydevd_pycharm.settrace('10.81.1.111', port=50000, stdoutToServer=True, stderrToServer=True)
  run_from_feed_cache_base(
      products,
      replayer.on_book_reset,
      start_time,
      end_time,
      fs_proto.feed_machine,
      fs_proto.worker_ids.split(','),
      write_feed_cache_if_none=fs_proto.write_feed_cache_if_none,
      feed_cache_dir=fs_proto.feed_cache_dir,
      api_override=fs_proto.api_override)
  pass


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  app.run(main)
